feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext by snehalv2002 · Pull Request #5052 · AI-Hypercomputer/maxtext

snehalv2002 · 2026-08-28T22:28:43Z

Description

Enhances MaxText's standalone checkpoint conversion tool (src/maxtext/checkpoint_conversion/to_maxtext.py) to directly ingest, serialize, and store native 8-bit float weights (float8_e4m3fn) and companion scale tensors into Orbax format without upcasting.

Motivation & Context

Previously, to_maxtext.py only supported save_dtype="bfloat16" and save_dtype="float32". Converting pre-quantized FP8 Hugging Face checkpoints forced CPU upcasting to BF16, doubling memory usage and storage footprint. Direct ingestion allows converting multi-gigabyte/terabyte FP8 checkpoints without memory bloat.

Key Changes

Direct PyTorch $\to$ JAX FP8 Bridging:
- Reinterprets PyTorch torch.float8_e4m3fn tensors as uint8 views before bridging to ml_dtypes.float8_e4m3fn in NumPy memory, bypassing standard NumPy float8 limitations.
- Updated LazyHFLoader to use safetensors with framework="pt" (when PyTorch is available) so native FP8 tensors can be deserialized from safetensors.
- Added ml_dtypes.float8_e4m3fn handling to LazyTensor dtype initialization.
Scale Key Fallback Resolution:
- Added resolve_scale_key() to resolve .weight_scale to .weight_scale_inv fallbacks in both LazyHFLoader and _eager_getter, supporting checkpoints with inverse scale naming conventions (e.g. NeuralMagic / vLLM).
Source Precision Preservation:
- When --save_dtype="float8_e4m3fn", _convert_tensor_to_numpy only converts FP8 weights to ml_dtypes.float8_e4m3fn. Companion scale tensors and unquantized modules (e.g., embeddings, norm scales, router gates) retain their source precision (BF16 or FP32).
- In convert_lora_to_maxtext_adapter, adapter weights remain bfloat16 when save_dtype="float8_e4m3fn".
1D Tensor Stacking & Slicing Safety (tensor_handling.py, utils.py, to_maxtext.py):
- In _build_single_axis_stacked_tensor (tensor_handling.py and to_maxtext.py), guarded stacking logic when scan_layers=True: if len(target_shape) <= param_scan_axis (such as 1D per-layer scalar scale tensors like Llama 3.1 FP8 kernel_scale), tensors are stacked along axis 0 instead of attempting an out-of-bounds stack along axis 1 (param_scan_axis).
- In process_maxtext_param (utils.py), added corresponding logic to slice along axis 0 for 1D parameters during MaxText parameter export.
CLI Normalization & Argument Extraction:
- Added _extract_conversion_args() and CLI argv preprocessing to allow conversion arguments (save_dtype, hf_model_path, lazy_load_tensors, etc.) to be supplied using standard MaxText key=value syntax without -- prefixes, separating them cleanly from pyconfig arguments.
- Expanded --save_dtype choices to include float8_e4m3fn.

Part 2 of 5 in the FP8 Weight-Only Dynamic Dequantization series (depends on #5053).

Tests

Automated Unit Tests

Added unit test suite in tests/unit/to_maxtext_test.py:

ResolveScaleKeyTest: Verifies exact key matching, .weight_scale $\to$ .weight_scale_inv fallback, and unmatched key handling.
ConvertTensorToNumpyTest: Verifies torch.float8_e4m3fn $\to$ ml_dtypes.float8_e4m3fn zero-copy view, FP8 $\to$ FP32 conversion, BF16 conversion, and NumPy passthrough.
ExtractConversionArgsTest: Verifies extraction and separation of conversion flags from mixed pyconfig CLI arguments.

End-to-End Verification

Tested end-to-end checkpoint conversion of Hugging Face FP8 model (neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8):

python -m maxtext.checkpoint_conversion.to_maxtext \
  model_name=llama3.1-8b-fp8 \
  hf_model_path=neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8 \
  base_output_directory=/tmp/maxtext_llama3.1_8b_fp8 \
  save_dtype=float8_e4m3fn \
  --lazy_load_tensors=True

Verified that linear weights are stored in float8_e4m3fn and scale tensors in float32.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

gemini-code-assist

Code Review

This pull request introduces support for FP8 weight-only storage and dynamic dequantization in MaxText, including changes to checkpoint conversion, linear layers, and decoder layers to support FP8 weights alongside scalar, per-channel, and block-wise scales. Feedback on these changes includes addressing a potential lossy double conversion of float16/bfloat16 tensors during checkpoint loading, extracting duplicated scale key resolution logic in to_maxtext.py into a helper function, and removing the redundant local _is_fp8_dtype helper in linears.py in favor of the existing utility in common_types.py.

gemini-code-assist · 2026-08-28T22:31:38Z

+          elif t.dtype == torch.bfloat16:
+            return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)
+          elif t.dtype == torch.float16:
+            return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)


Converting torch.float16 and torch.bfloat16 unconditionally to ml_dtypes.bfloat16 inside get_tensor can lead to a lossy double conversion when save_dtype is configured as float32 (i.e., float16 -> bfloat16 -> float32).

Since LazyTensor.__array__ already handles casting the retrieved numpy array to the target save_dtype (via arr.astype(dtype)), get_tensor should simply return the closest native numpy representation (e.g., float32 for bfloat16 to avoid PyTorch conversion issues, and native float16 for float16).

Suggested change

elif t.dtype == torch.bfloat16:

return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)

elif t.dtype == torch.float16:

return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)

elif t.dtype == torch.bfloat16:

return t.to(torch.float32).numpy()

elif t.dtype == torch.float16:

return t.numpy()

gemini-code-assist · 2026-08-28T22:31:38Z

+    if shard_name is None:
+      # Check fallback for .weight_scale vs .scale and inverse scales
+      if resolved_key.endswith(".weight_scale"):
+        for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+          alt_key = resolved_key[:-len(".weight_scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break
+      elif resolved_key.endswith(".scale"):
+        for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+          alt_key = resolved_key[:-len(".scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break


The fallback logic for resolving .weight_scale vs .scale and inverse scales is duplicated three times in this file (here, in get_tensor's second block, and in _eager_getter).

To improve maintainability and reduce redundancy, consider extracting this logic into a helper function at the module level:

def resolve_scale_key(key: str, container) -> str: """Resolves fallback keys for weight scales and inverse scales.""" if key in container: return key if key.endswith(".weight_scale"): for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]: alt_key = key[:-len(".weight_scale")] + suffix if alt_key in container: return alt_key elif key.endswith(".scale"): for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]: alt_key = key[:-len(".scale")] + suffix if alt_key in container: return alt_key return key

You can then simplify this block to:

resolved_key = resolve_scale_key(key, self.shard_map) shard_name = self.shard_map.get(resolved_key)

gemini-code-assist · 2026-08-28T22:31:38Z

+def _is_fp8_dtype(dtype: Any) -> bool:
+  """Checks whether a dtype is an FP8 data type."""
+  if dtype is None:
+    return False
+  try:
+    canon_dtype = _canonicalize_dtype(dtype)
+  except (TypeError, ValueError):
+    return False
+
+  fp8_types = [jnp.float8_e4m3fn, jnp.float8_e5m2]
+  for attr in ("float8_e4m3fnuz", "float8_e5m2fnuz", "float8_e4m3b11fnuz"):
+    if hasattr(jnp, attr):
+      fp8_types.append(getattr(jnp, attr))
+  return canon_dtype in fp8_types


The _is_fp8_dtype helper function is duplicated here. MaxText already has a standard is_fp8_dtype utility defined in maxtext.common.common_types.

You can import is_fp8_dtype from maxtext.common.common_types and use it directly instead of redefining _is_fp8_dtype locally.

snehalv2002 · 2026-09-09T20:14:06Z

+
+if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
+  os.environ["HF_HOME"] = "/dev/shm/hf_cache"
+


remove this it was only for local development

snehalv2002 · 2026-09-09T20:45:34Z

  checkpointing.wait_until_finished(checkpoint_manager)

  max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min")
+


delete this new line addition

snehalv2002 · 2026-09-09T20:47:10Z

+      with safe_open(local_path, framework=framework, device="cpu") as f:
+        final_key = resolve_scale_key(resolved_key, f.keys())
+        t = f.get_tensor(final_key)
+        if torch is not None and isinstance(t, torch.Tensor):
+          if hasattr(torch, "float8_e4m3fn") and t.dtype == torch.float8_e4m3fn:
+            return t.view(torch.uint8).numpy().view(ml_dtypes.float8_e4m3fn)
+          elif hasattr(torch, "float8_e5m2") and t.dtype == torch.float8_e5m2:
+            return t.view(torch.uint8).numpy().view(ml_dtypes.float8_e5m2)
+          elif t.dtype == torch.bfloat16:
+            if self.save_dtype in ("float32", DType.FLOAT32):
+              return t.to(torch.float32).numpy()
+            return t.view(torch.int16).numpy().view(ml_dtypes.bfloat16)
+          elif t.dtype == torch.float16:
+            if self.save_dtype in ("float32", DType.FLOAT32):
+              return t.to(torch.float32).numpy()
+            elif self.save_dtype in ("bfloat16", DType.BFLOAT16):
+              return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)
+            return t.numpy()
+          else:
+            return t.numpy()
+        return t


can we simplify this somehow?

gobbleturk · 2026-09-10T17:37:50Z

  if config.scan_layers:
    # If it's a standard scanned layer, we use the configured param_scan_axis.
-    axis_to_stack = config.param_scan_axis
+    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0


can you include a comment here?

My guess is this is used for 1D scales?

Yes for 1D scales. Added a comment to specify this

codecov · 2026-09-10T21:27:55Z

Codecov Report

❌ Patch coverage is 6.12245% with 92 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/checkpoint_conversion/to_maxtext.py	5.26%	89 Missing and 1 partial ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

…stion in to_maxtext Extends the MaxText HuggingFace checkpoint conversion pipeline to directly ingest FP8 weight tensors and companion scale tensors. Key Changes: - to_maxtext: Support float8_e4m3fn and float8_e5m2 save dtypes in _convert_tensor_to_numpy and _eager_getter, preserving each tensor's source precision (companion scales keep their source dtype; unquantized modules stay bfloat16). - to_maxtext: Support scale key suffix resolution (.weight_scale, .scale, .scale_inv, .weight_scale_inv) in LazyHFLoader and _eager_getter. - tensor_handling & to_maxtext: Handle 1D per-layer scalar scale stacking along axis 0 when scan_layers=True. - utils: Support slicing 1D MaxText parameters along axis 0 during checkpoint export. - tests: Add unit tests verifying FP8 tensor casting and scale suffix resolution in hf_checkpoint_conversion_test.py.

gemini-code-assist Bot reviewed Aug 28, 2026

View reviewed changes

snehalv2002 changed the title ~~Pr/fp8 to maxtext~~ feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext Aug 28, 2026

This was referenced Aug 28, 2026

feat(checkpointing): add dequantize-on-load parameter restoration #5051

Merged

feat(models): add llama3.1-8b-fp8 model config and HuggingFace mapping #5054

Draft

test(fp8): add PyTorch reference parity suite and logit checker metrics #5055

Draft

snehalv2002 changed the base branch from main to pr/fp8-dequant-engine August 28, 2026 22:52

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 7687dfa to f063fd7 Compare September 2, 2026 17:56

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 56098db to 3d13d99 Compare September 2, 2026 17:56

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from f063fd7 to 8e544e6 Compare September 2, 2026 21:27

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 3d13d99 to 4e38337 Compare September 2, 2026 21:27

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 8e544e6 to f74726d Compare September 2, 2026 22:25

snehalv2002 force-pushed the pr/fp8-to-maxtext branch 6 times, most recently from cb0ebb5 to f54bfa9 Compare September 8, 2026 20:43

snehalv2002 commented Sep 9, 2026

View reviewed changes

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 9456383 to c3e0942 Compare September 9, 2026 20:32

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from f54bfa9 to 44cc3a7 Compare September 9, 2026 20:43

snehalv2002 commented Sep 9, 2026

View reviewed changes

snehalv2002 marked this pull request as ready for review September 9, 2026 22:09

snehalv2002 requested review from RissyRan, bvandermoon, gobbleturk, hengtaoguo, jiangjy1982, parambole and shuningjin as code owners September 9, 2026 22:09

gobbleturk reviewed Sep 10, 2026

View reviewed changes

gobbleturk approved these changes Sep 10, 2026

View reviewed changes

Shuwen-Fang self-requested a review September 10, 2026 20:46

Shuwen-Fang approved these changes Sep 10, 2026

View reviewed changes

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 2ae774c to dce35b2 Compare September 10, 2026 21:04

snehalv2002 force-pushed the pr/fp8-to-maxtext branch 2 times, most recently from dbf4b02 to fde4737 Compare September 10, 2026 21:10

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from dce35b2 to c8f282a Compare September 10, 2026 21:10

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from fde4737 to ed14346 Compare September 11, 2026 00:39

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from c8f282a to 82dc9d4 Compare September 11, 2026 17:47

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from ed14346 to 205de7a Compare September 11, 2026 17:47

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 82dc9d4 to b74c55f Compare September 11, 2026 18:22

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 205de7a to 41f7756 Compare September 11, 2026 18:22

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from b74c55f to 279e76e Compare September 11, 2026 21:24

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 41f7756 to ebc3f73 Compare September 11, 2026 21:24

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 279e76e to 02eac58 Compare September 11, 2026 21:27

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from ebc3f73 to cb11b27 Compare September 11, 2026 21:27

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 02eac58 to 6563931 Compare September 12, 2026 00:13

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from cb11b27 to 5e61891 Compare September 12, 2026 00:14

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 6563931 to 921921e Compare September 12, 2026 00:16

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 5e61891 to f3a466f Compare September 12, 2026 00:16

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 921921e to 219da46 Compare September 12, 2026 04:13

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from f3a466f to e4f3ad5 Compare September 12, 2026 04:14

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from 219da46 to a19e63c Compare September 12, 2026 10:59

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from e4f3ad5 to 69bc314 Compare September 12, 2026 11:00

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from a19e63c to d32cf49 Compare September 12, 2026 20:35

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 69bc314 to c84e000 Compare September 12, 2026 20:35

snehalv2002 force-pushed the pr/fp8-dequant-engine branch from d32cf49 to 1bb4531 Compare September 14, 2026 19:40

Navigation Menu

feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext - #5052

feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext#5052
copybara-service[bot] merged 1 commit into
mainfrom
pr/fp8-to-maxtext

snehalv2002 commented Aug 28, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

snehalv2002 Sep 9, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

snehalv2002 Sep 9, 2026 •

edited

Loading

Uh oh!

snehalv2002 Sep 9, 2026

Uh oh!

snehalv2002 Sep 9, 2026

Uh oh!

gobbleturk Sep 10, 2026

Uh oh!

snehalv2002 Sep 10, 2026

Uh oh!

codecov Bot commented Sep 10, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants


		if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
		os.environ["HF_HOME"] = "/dev/shm/hf_cache"

		checkpointing.wait_until_finished(checkpoint_manager)

		max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min")

Conversation

snehalv2002 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation & Context

Key Changes

Tests

Automated Unit Tests

End-to-End Verification

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gobbleturk Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

snehalv2002 commented Aug 28, 2026 •

edited

Loading

snehalv2002 Sep 9, 2026 •

edited

Loading

codecov Bot commented Sep 10, 2026 •

edited

Loading